package com.yonghui.springsecurityoauth2demo.config;

import com.yonghui.springsecurityoauth2demo.pojo.LoginUser;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;

import java.util.HashMap;
import java.util.Map;

/**
 * @author liyonghui
 * @description: JWT内容增强
 * @date 2021/12/5 15:58
 */
public class JwtTokenEnhancer implements TokenEnhancer {

    // @Override
    // public OAuth2AccessToken enhance(OAuth2AccessToken oAuth2AccessToken, OAuth2Authentication oAuth2Authentication) {
    //     Map<String, Object> objectObjectHashMap = new HashMap<>();
    //     objectObjectHashMap.put("enhance", "enhance info");
    //     objectObjectHashMap.put("ceshi", "张三");
    //     ((DefaultOAuth2AccessToken) oAuth2AccessToken).setAdditionalInformation(objectObjectHashMap);
    //     return oAuth2AccessToken;
    // }

    @Override
    public OAuth2AccessToken enhance(OAuth2AccessToken accessToken, OAuth2Authentication oAuth2Authentication) {
        if (accessToken instanceof DefaultOAuth2AccessToken) {
            DefaultOAuth2AccessToken token = (DefaultOAuth2AccessToken) accessToken;
            token.setRefreshToken(null); // 移除refresh_token
        }


        Map<String, Object> additionalInfo = new HashMap<>();
        // 这里可以根据需要添加更多信息
        //additionalInfo.put("organization", oAuth2Authentication.getName() + " from my Organization");
        // 如果用户信息中包含特定的属性，可以这样添加
        LoginUser principal = (LoginUser)oAuth2Authentication.getPrincipal();
        additionalInfo.put("userId", principal.getUserId());  // 假设 principal 中包含了用户ID

        // 可以添加更多的自定义属性
        ((DefaultOAuth2AccessToken) accessToken).setAdditionalInformation(additionalInfo);
        return accessToken;
    }

}
